"""
原始版权声明
Copyright (c) 2025 Mr_寒路
此文件最初源自 https://gitee.com/mr-hanlu/llm-cli-lite/tree/master/llm_cli_lite ，遵循MIT许可证。
"""
from openai import OpenAI
import base64

class ChatCompletionClient:
    def __init__(self, base_url: str, api_key: str, model_name: str, system: str):
        self.client = OpenAI(base_url=base_url, api_key=api_key)
        self.model_name = model_name
        self.max_tokens = 4096
        self.temperature = 0.8
        self.top_k = 20
        self.top_p = 0.9

        self.system = [{
            "role": "system",
            "content": system,
        }]

    def get_models(self) -> list:
        # 获取可用的模型列表
        models = self.client.models.list()
        return [model.id for model in models.data]

    def create_chat_completion(self, prompt: str, image_path: str = "", stream: bool = False) -> str:
        # 纯文本问答
        message = self.system + [{"role": "user", "content": prompt}]
        # 创建聊天完成
        out = self.client.chat.completions.create(
            model=self.model_name,
            messages=message,
            temperature=self.temperature,
            top_p=self.top_p,
            stream=stream
        )
        print("")
        answer = ""
        for chunk in out:
            answer_chunk = (chunk.choices[0].delta.content or "")
            answer += answer_chunk
            print(answer_chunk, end="")
        print() 
        print() 

        return answer

    def set_model(self, model_name: str):
        # 设置模型名称
        self.model_name = model_name

    def set_params(self, max_tokens: int = 200, temperature: float = 0.5, top_p: float = 0.7):
        # 设置请求参数
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.top_p = top_p

